from PyQt5 import uic,QtCore,QtGui,QtWidgets
from PyQt5.QtWidgets import QWidget,QApplication

import cv2
import torch
import sys
import os


class MyWindow(QWidget):
    

    def __init__(self):
        super().__init__()
        self.init_ui()
        

    def init_ui(self):
        
        self.timer_camera = QtCore.QTimer()  # 定义定时器，用于控制显示视频的帧率
        
        
        #加载UI界面
        self.ui = uic.loadUi("./Safetyhelmet_detect_v2.0.ui")
        # print(self.ui.__dict__)  # 查看ui文件中有哪些控件

        # 设置标题
        self.ui.setWindowTitle("安全帽检测系统 V2.0")
        # 设置图标
        self.ui.setWindowIcon(QtGui.QIcon('./data/log.png'))

        # 提取要操作的控件
        self.show_result_label = self.ui.show_result_label  # 结果展示框
        self.models_comboBox = self.ui.models_comboBox  # 模型选择框
        self.start_detect_btn = self.ui.start_detect_btn  # 开始检测按钮
        self.images_comboBox = self.ui.images_comboBox  # 照片路径
        self.show_log_label = self.ui.show_log_label
        self.openclose_camera_btn = self.ui.openclose_camera_btn
        self.load_imagefile_btn = self.ui.load_imagefile_btn
        self.load_model_btn =  self.ui.load_model_btn
        self.load_model_label  = self.ui.load_model_label
        self.detect_result_label = self.ui.detect_result_label
        
        # self.load_model()
        # self.model = torch.hub.load('./yolov5', 'custom', path='./models/model_safetyhelmet_detect_yolo_s.pt',source='local')  # local repo
        # self.model.conf = 0.3
        # self.load_model_label.setText("模型已加载！")
        # self.load_model_label.setAlignment(QtCore.Qt.AlignCenter)#居中


        self.slot_init()
        self.loadmodel_comboBox()
        self.loadimage_comboBox()

        
        self.show_logimg("./data/log1.png",self.show_log_label)
        self.show_logimg("./data/log2.png",self.show_result_label)




    def slot_init(self):

        self.start_detect_btn.clicked.connect(self.start_detect_img)  # 开始检测按钮绑定信号与槽函数
        self.openclose_camera_btn.clicked.connect(self.button_open_camera_clicked)  # 若该按键被点击，则调用button_open_camera_clicked()
        self.timer_camera.timeout.connect(self.show_camera)  # 若定时器结束，则调用show_camera()
        self.load_imagefile_btn.clicked.connect(self.loadmodel_comboBox)#更新检测照片
        self.load_model_btn.clicked.connect(self.load_model)


    def loadmodel_comboBox(self):

        models_path = ("./models")
        models_list = os.listdir(models_path)
        self.models_comboBox.clear()
        #添加到下拉框预选
        self.models_comboBox.addItems([models_list[i] for i in range(len(models_list))])

    def loadimage_comboBox(self):

        self.images_path = ("./detect_images")
        images_list = os.listdir(self.images_path)
        self.images_comboBox.clear()
        #添加到下拉框预选
        self.images_comboBox.addItems([images_list[i] for i in range(len(images_list))])
   
   
    def load_model(self):
        
        self.load_model_label.setText("模型加载中ing...")
        self.load_model_label.setStyleSheet("color:red")
        self.load_model_label.setAlignment(QtCore.Qt.AlignCenter)#居中
        self.load_model_label.repaint()#更新画面
        model_name = self.models_comboBox.currentText()
        
            
        if model_name == "model_safetyhelmet_detect_yolo_s.pt":
            self.model = torch.hub.load('./yolov5', 'custom', path='./models/model_safetyhelmet_detect_yolo_s.pt',source='local')  # local repo
        if model_name == "model_safetyhelmet_detect_yolo_m.pt":
            self.model = torch.hub.load('./yolov5', 'custom', path='./models/model_safetyhelmet_detect_yolo_m.pt',source='local')  # local repo
        if model_name == "model_safetyhelmet_detect_yolo_n.pt":
            self.model = torch.hub.load('./yolov5', 'custom', path='./models/model_safetyhelmet_detect_yolo_n.pt',source='local')  # local repo
        if model_name == "model_safetyhelmet_detect_yolo_n6.pt":
            self.model = torch.hub.load('./yolov5', 'custom', path='./models/model_safetyhelmet_detect_yolo_n6.pt',source='local')  # local repo


        self.load_model_label.setText(model_name + "_" + "模型加载完毕！")
        self.load_model_label.setAlignment(QtCore.Qt.AlignCenter)#居中


    #打开摄像头按钮
    def button_open_camera_clicked(self):
        self.cap = cv2.VideoCapture()  # 视频流
        self.CAM_NUM = 0  # 为0时表示视频流来自笔记本内置摄像头
        if self.timer_camera.isActive() == False:  # 若定时器未启动
            flag = self.cap.open(0)  # 参数是0，表示打开笔记本的内置摄像头，参数是视频文件路径则打开视频
            if flag == False:  # flag表示open()成不成功
                msg = QtWidgets.QMessageBox.warning(self, 'warning', "请检查相机于电脑是否连接正确", buttons=QtWidgets.QMessageBox.Ok)
            else:
                self.timer_camera.start(30)  # 定时器开始计时30ms，结果是每过30ms从摄像头中取一帧显示
                self.openclose_camera_btn.setText('关闭相机')
        else:
            self.timer_camera.stop()  # 关闭定时器
            self.cap.release()  # 释放视频流
            # self.cap.release()
            self.detect_result_label.setText("")
            self.detect_result_label.setStyleSheet("font-size: 30px; color: red") 
            self.show_result_label.clear()  # 清空视频显示区域
            self.openclose_camera_btn.setText('打开相机')
            self.show_logimg("./data/log2.png",self.show_result_label)
 

    def show_camera(self):

        flag, self.image = self.cap.read()  # 从视频流中读取

        img = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)  # 视频色彩转换回RGB，这样才是现实的颜色

        self.detect_img(img)


    def start_detect_img(self):
        imgname = self.images_comboBox.currentText()
        img = cv2.imread(self.images_path + "\\" +imgname)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        self.detect_img(img)


    def detect_img(self,img):#开始检测并绘制检测结果到图像
        

        model_state = self.load_model_label.text().split('_')[-1]
        # print(model_state)
        if model_state == "模型加载完毕！":
            img = cv2.resize(img,(430,400))
            results = self.model(img)
            pd = results.pandas().xyxy[0]

            person_list = pd[pd['name']=='person'].to_numpy()
            vest_list = pd[pd['name']=='vest'].to_numpy()
            hat_list = pd[pd['name'].str.contains('helmet')].to_numpy()


            #绘制人体框
            for person in person_list:
                l,t,r,b = person[:4].astype('int')
                cv2.rectangle(img,(l,t),(r,b),(0,255,0),5)

            if len(person_list)== 0 and len(vest_list)== 0 and len(hat_list)== 0 and self.openclose_camera_btn.text() == '关闭相机':
                self.detect_result_label.setText("--带上安全帽,安全立刻到！--")
                self.detect_result_label.setAlignment(QtCore.Qt.AlignCenter)#居中
                self.detect_result_label.setStyleSheet("font-size: 30px; color: red; background-color: rgb(255, 251, 100)") 

            elif len(person_list) != len(hat_list)  and  len(person_list) != 0 :

                # cv2.putText(img,'Helmet wearing is not regulated!',(10,20),cv2.FONT_HERSHEY_SIMPLEX,1,(255,0,0),2)
                self.detect_result_label.setText("---请规范佩戴安全头盔！---")
                self.detect_result_label.setAlignment(QtCore.Qt.AlignCenter)#居中
                self.detect_result_label.setStyleSheet("font-size: 30px; color: red; background-color: rgb(255, 251, 100)") 

            # #绘制人体框
            # for person in person_list:
            #     l,t,r,b = person[:4].astype('int')
            #     cv2.rectangle(img,(l,t),(r,b),(255,0,0),5)

            elif len(person_list) == len(hat_list):
                # cv2.putText(img,'---Pass!---',(10,20),cv2.FONT_HERSHEY_SIMPLEX,1,(0,255,0),2)
                self.detect_result_label.setText("---PASS---")
                self.detect_result_label.setAlignment(QtCore.Qt.AlignCenter)#居中
                self.detect_result_label.setStyleSheet("font-size: 30px; color: green;background-color: rgb(255, 251, 100)") 


            elif self.openclose_camera_btn.text() == '关闭相机':
                self.detect_result_label.setText("---FAIL---")
           
            #绘制人体框
            for person in person_list:
                l,t,r,b = person[:4].astype('int')
                if self.detect_result_label.text() == "---PASS---":
                    cv2.rectangle(img,(l,t),(r,b),(0,255,0),5)
                else:
                    cv2.rectangle(img,(l,t),(r,b),(0,0,255),5)

            #绘制帽子框
            for hat in hat_list:
                l,t,r,b = hat[:4].astype('int')
                if self.detect_result_label.text() == "---PASS---":
                    cv2.rectangle(img,(l,t),(r,b),(0,255,0),5)
                else:
                    cv2.rectangle(img,(l,t),(r,b),(0,0,255),5)


            self.show_img_to_label(img)#展示检测结果到展示框
        else:
            self.show_img_to_label(img)#展示检测结果到展示框
            self.detect_result_label.setText("--请先加载检测模型!--")
            self.detect_result_label.setAlignment(QtCore.Qt.AlignCenter)#居中
            self.detect_result_label.setStyleSheet("font-size: 30px; color: red;background-color: rgb(255, 251, 100)") 

        
    def show_img_to_label(self,img):

        show = cv2.resize(img,(500,400))
        showImage = QtGui.QImage(show.data, show.shape[1], show.shape[0],QtGui.QImage.Format_RGB888)  # 把读取到的视频数据变成QImage形式
        self.show_result_label.setPixmap(QtGui.QPixmap.fromImage(showImage))  # 往显示视频的Label里 显示QImage

    #初始化界面显示校徽log
    def show_logimg(self,imgname,labelname):

        pix = QtGui.QPixmap(imgname)
        labelname.setPixmap(pix)
        labelname.setScaledContents(True)

if __name__ == '__main__':

    app = QApplication(sys.argv)

    w = MyWindow()

    # 展示窗口
    w.ui.show()

    app.exec()